
module IRGenerator = struct
  open IR  (* 打开 IR 模块以引用其类型定义 *)
  open Ast  (* 打开 Ast 模块以引用其类型定义 *)
  let current_break_label = ref ""
  let current_continue_label = ref ""
  let counter = ref 0
  let label_counter = ref 0
  let new_temp () = (*生成临时变量名*)
    let t = !counter in
    counter := t + 1;
    IR_Temp t
  let new_label prefix =(*生成标签名*)
    let l = !label_counter in
    label_counter := l + 1;
    prefix ^ "_" ^ string_of_int l

  let rec expr_to_ir expr = (*表达式转化，函数结果为(当前语句赋值的临时变量名:IR_Temp _，当前语句执行的操作列表:ir_instruction list)*)
    match expr with
    | Num n -> (IR_Const n, [])
    | ID id -> (IR_Var id, [])
    | Unary (op, e) ->
        let (operand, instrs) = expr_to_ir e in
        let temp = new_temp () in
        (temp, instrs @ [IR_Unary (op, operand, temp)])
    | Binary (op, e1, e2) ->
        let (op1, instrs1) = expr_to_ir e1 in
        let (op2, instrs2) = expr_to_ir e2 in
        let temp = new_temp () in
        (temp, instrs1 @ instrs2 @ [IR_Binary (op, op1, op2, temp)])
    | Call (fname, args) ->
        let (args_ir, args_instrs) = List.fold_left (fun (args, instrs) arg ->
          let (arg_ir, arg_instrs) = expr_to_ir arg in
          (args @ [arg_ir], instrs @ arg_instrs)
        ) ([], []) args in
        let temp = new_temp () in
        (temp, args_instrs @ [IR_Call (fname, args_ir, Some temp)])
    | Paren e -> expr_to_ir e

  let rec stmt_to_ir stmt = (*语句转化函数，函数结果为ir_instruction list*)
    match stmt with
    | Block stmts ->
        List.fold_left (fun instrs s ->
          let new_instrs = stmt_to_ir s in
          instrs @ new_instrs
        ) [] stmts
    | Empty -> []
    | ExprStmt e -> 
        let (_, instrs) = expr_to_ir e in
        instrs
    | Assign (id, e) ->
        let (operand, instrs) = expr_to_ir e in
        instrs @ [IR_Assign (IR_Var id, operand)]
    | Declare (id, e) ->
        let (operand, instrs) = expr_to_ir e in
        instrs @ [IR_Assign (IR_Var id, operand)]
    | If (cond, then_stmt, else_stmt) ->
        let (cond_operand, cond_instrs) = expr_to_ir cond in
        let then_label = new_label "then" in
        let else_label = new_label "else" in
        let end_label = new_label "endif" in
        let then_instrs = stmt_to_ir then_stmt in
        let else_instrs = match else_stmt with
          | Some s -> stmt_to_ir s
          | None -> [] in
        
        cond_instrs @
        [IR_If (cond_operand, then_label);
         IR_Goto else_label;
         IR_Label then_label] @
        then_instrs @
        [IR_Goto end_label;
         IR_Label else_label] @
        else_instrs @
        [IR_Label end_label]
    | While (cond, body) ->
        let cond_label = new_label "while_cond" in
        let body_label = new_label "while_body" in
        let end_label = new_label "while_end" in
        let (cond_operand, cond_instrs) = expr_to_ir cond in
        let body_instrs = stmt_to_ir body in

        current_break_label := end_label; (*更新break语句跳转地址*)
        current_continue_label := cond_label; (*更新continue语句跳转地址*)

        [IR_Label cond_label] @
        cond_instrs @
        [IR_If (cond_operand, body_label);
         IR_Goto end_label;
         IR_Label body_label] @
        body_instrs @
        [IR_Goto cond_label;
         IR_Label end_label]
    | Break -> [IR_Goto !current_break_label]
    | Continue -> [IR_Goto !current_continue_label]
    | Return (Some e) ->
        let (operand, instrs) = expr_to_ir e in
        instrs @ [IR_Return (Some operand)]
    | Return None -> [IR_Return None]

  let func_def_to_ir func = (*函数转化函数，函数结果为ir_function*)
    let param_names = List.map (fun p -> p.pid) func.params in
    let body_ir = stmt_to_ir func.body in
    
let blocks = 
  [{
    label = Some "entry";
    instructions = body_ir @ [IR_Return None];
    locals = []; (* 确保 locals 字段存在 *)
    terminator = IR_Return None (* 实际应该从body_ir中提取 *)
  }] in
    
    {
      name = func.name;
      params = param_names;
      locals = param_names; 
      blocks = blocks;
    }

  let program_to_ir prog =
    List.map func_def_to_ir prog
end